import torch
import torch.nn as nn
from torch.optim import SGD


device = torch. device("cuda" if torch.cuda.is_available() else "cpu")

x = torch.rand([500, 1]).to(device)
y_true = 3 * x + 0.8


class MyLinear(nn.Module):
    def __init__(self):
        super(MyLinear,self).__init__()
        self.linear = nn.Linear(1,1)


    def forward(self,x):
        out = self.linear(x)
        return out


my_linear =MyLinear().to(device)
optimizer = SGD(my_linear.parameters(),0.005)
loss_fn = nn.MSELoss()


for i in range(30000):
    y_predict = my_linear(x)
    loss = loss_fn(y_predict,y_true)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if i % 2000 == 0:
        print(loss.item(),list(my_linear.parameters()))




